{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import re\n",
      "from regex import D\n",
      "\n",
      "import rope.base.project\n",
      "from loguru import logger\n",
      "from rope.refactor.extract import ExtractMethod\n",
      "\n",
      "from sweepai.agents.name_agent import NameBot\n",
      "from sweepai.config.server import DEFAULT_GPT4_32K_MODEL, DEFAULT_GPT35_MODEL\n",
      "from sweepai.core.chat import ChatGPT\n",
      "from sweepai.core.entities import Message\n",
      "from sweepai.core.update_prompts import (\n",
      "    extract_snippets_system_prompt,\n",
      "    extract_snippets_user_prompt,\n",
      ")\n",
      "from sweepai.utils.github_utils import ClonedRepo\n",
      "from sweepai.utils.jedi_utils import (\n",
      "    get_all_defined_functions,\n",
      "    get_references_from_defined_function,\n",
      "    setup_jedi_for_file,\n",
      ")\n",
      "from sweepai.utils.refactor_utils import get_refactor_snippets\n",
      "from sweepai.utils.search_and_replace import find_best_match\n",
      "\n",
      "APOSTROPHE_MARKER = \"__APOSTROPHE__\"\n",
      "PERCENT_FORMAT_MARKER = \"__PERCENT_FORMAT__\"\n",
      "\n",
      "\n",
      "def serialize(text: str):\n",
      "    # Replace \"'{var}'\" with \"__APOSTROPHE__{var}__APOSTROPHE__\"\n",
      "    text = re.sub(\n",
      "        r\"'{([^'}]*?)}'\", f\"{APOSTROPHE_MARKER}{{\\\\1}}{APOSTROPHE_MARKER}\", text\n",
      "    )\n",
      "    # Replace \"%s\" with \"__PERCENT_FORMAT__\"\n",
      "    text = re.sub(r\"%\\((.*?)\\)s\", f\"{PERCENT_FORMAT_MARKER}{{\\\\1}}\", text)\n",
      "    return text\n",
      "\n",
      "\n",
      "def deserialize(text: str):\n",
      "    text = re.sub(f\"{APOSTROPHE_MARKER}{{(.*?)}}{APOSTROPHE_MARKER}\", \"'{\\\\1}'\", text)\n",
      "    text = re.sub(f\"{PERCENT_FORMAT_MARKER}{{(.*?)}}\", \"%(\\\\1)s\", text)\n",
      "    return text\n",
      "\n",
      "def count_plus_minus_in_diff(description):\n",
      "    plus_count = sum([1 for line in description.split(\"\\n\") if line.startswith(\"+\")])\n",
      "    minus_count = sum([1 for line in description.split(\"\\n\") if line.startswith(\"-\")])\n",
      "    return plus_count, minus_count\n",
      "\n",
      "def extract_method(\n",
      "    snippet,\n",
      "    file_path,\n",
      "    method_name,\n",
      "    project_name,\n",
      "):\n",
      "    project = rope.base.project.Project(project_name)\n",
      "\n",
      "    resource = project.get_resource(file_path)\n",
      "    contents = resource.read()\n",
      "    serialized_contents = serialize(contents)\n",
      "    resource.write(serialized_contents)\n",
      "\n",
      "    serialized_snippet = serialize(snippet)\n",
      "    start, end = serialized_contents.find(serialized_snippet), serialized_contents.find(\n",
      "        serialized_snippet\n",
      "    ) + len(serialized_snippet)\n",
      "\n",
      "    try:\n",
      "        import pdb; pdb.set_trace()\n",
      "        extractor = ExtractMethod(project, resource, start, end)\n",
      "        change_set = extractor.get_changes(method_name, similar=True)\n",
      "\n",
      "        for change in change_set.changes:\n",
      "            if change.old_contents is not None:\n",
      "                change.old_contents = deserialize(change.old_contents)\n",
      "            else:\n",
      "                change.old_contents = deserialize(change.resource.read())\n",
      "            change.new_contents = deserialize(change.new_contents)\n",
      "\n",
      "        # adding this because the change might not replace old code.\n",
      "        _, subtracted_lines = count_plus_minus_in_diff(change_set.get_description())\n",
      "        if (\n",
      "            subtracted_lines <= 3\n",
      "        ):\n",
      "            logger.info(\"Change doesn't remove code, skipping\")\n",
      "            return contents, []\n",
      "        for change in change_set.changes:\n",
      "            change.do()\n",
      "\n",
      "        result = deserialize(resource.read())\n",
      "        resource.write(result)\n",
      "        return result, change_set\n",
      "    except Exception as e:\n",
      "        logger.error(f\"An error occurred: {e}\")\n",
      "        resource.write(contents)\n",
      "        return contents, []\n",
      "\n",
      "\n",
      "class RefactorBot(ChatGPT):\n",
      "    def refactor_snippets(\n",
      "        self,\n",
      "        additional_messages: list[Message] = [],\n",
      "        snippets_str=\"\",\n",
      "        file_path: str = \"\",\n",
      "        update_snippets_code: str = \"\",\n",
      "        request=\"\",\n",
      "        changes_made=\"\",\n",
      "        cloned_repo: ClonedRepo = None,\n",
      "        **kwargs,\n",
      "    ):\n",
      "        self.model = (\n",
      "            DEFAULT_GPT4_32K_MODEL\n",
      "            if (self.chat_logger and self.chat_logger.is_paying_user())\n",
      "            else DEFAULT_GPT35_MODEL\n",
      "        )\n",
      "\n",
      "        # first perform manual refactoring step\n",
      "        script, tree = setup_jedi_for_file(\n",
      "            project_dir=cloned_repo.repo_dir,\n",
      "            file_full_path=f\"{cloned_repo.repo_dir}/{file_path}\",\n",
      "        )\n",
      "\n",
      "        all_defined_functions = get_all_defined_functions(script=script, tree=tree)\n",
      "        initial_file_contents = cloned_repo.get_file_contents(file_path=file_path)\n",
      "        heuristic_based_extractions = get_refactor_snippets(initial_file_contents, {}) # check heuristics\n",
      "        if len(heuristic_based_extractions) > 0:\n",
      "            # some duplicated code here\n",
      "            deduped_exact_matches = heuristic_based_extractions  # already deduped\n",
      "            new_function_names = []\n",
      "            existing_names = \", \".join(\n",
      "                [def_fn.name.strip(\"'\") for def_fn in all_defined_functions]\n",
      "            )\n",
      "            offset = 5\n",
      "            for idx in range(0, len(deduped_exact_matches), offset):\n",
      "                num_snippets = min(len(deduped_exact_matches), idx + offset) - idx\n",
      "                formatted_snippets = \"\\n\".join(\n",
      "                    [\n",
      "                        f\"<function_to_name>\\n{snippet}\\n</function_to_name>\"\n",
      "                        for snippet in deduped_exact_matches[idx:idx+num_snippets]\n",
      "                    ]\n",
      "                )\n",
      "                new_function_names.extend(NameBot(chat_logger=self.chat_logger).name_functions(\n",
      "                    old_code=cloned_repo.get_file_contents(file_path),\n",
      "                    snippets=formatted_snippets,\n",
      "                    existing_names=existing_names,\n",
      "                    count=num_snippets,\n",
      "                ))\n",
      "            for idx, extracted_original_code in enumerate(deduped_exact_matches):\n",
      "                if idx >= len(new_function_names):\n",
      "                    break\n",
      "                new_code, _ = extract_method(\n",
      "                    extracted_original_code,\n",
      "                    file_path,\n",
      "                    new_function_names[idx],\n",
      "                    project_name=cloned_repo.repo_dir,\n",
      "                )\n",
      "\n",
      "        self.messages = [\n",
      "            Message(\n",
      "                role=\"system\",\n",
      "                content=extract_snippets_system_prompt,\n",
      "                key=\"system\",\n",
      "            )\n",
      "        ]\n",
      "        self.messages.extend(additional_messages)\n",
      "        new_code = None\n",
      "        extracted_exact_matches = []\n",
      "        new_function_names = []\n",
      "        for fn_def in all_defined_functions:\n",
      "            full_file_code = cloned_repo.get_file_contents(file_path)\n",
      "            script, tree = setup_jedi_for_file(\n",
      "                project_dir=cloned_repo.repo_dir,\n",
      "                file_full_path=f\"{cloned_repo.repo_dir}/{file_path}\",\n",
      "            )\n",
      "            function_and_reference = get_references_from_defined_function(\n",
      "                fn_def,\n",
      "                script,\n",
      "                tree,\n",
      "                f\"{cloned_repo.repo_dir}/{file_path}\",\n",
      "                full_file_code,\n",
      "            )\n",
      "            if function_and_reference.function_code.count(\"\\n\") < 20:\n",
      "                continue\n",
      "            # everything below must operate in a loop\n",
      "            recent_file_contents = cloned_repo.get_file_contents(file_path=file_path)\n",
      "            code = function_and_reference.serialize(tag=\"function_to_refactor\")\n",
      "            extract_response = self.chat(\n",
      "                extract_snippets_user_prompt.format(\n",
      "                    code=code,\n",
      "                    file_path=file_path,\n",
      "                    snippets=snippets_str,\n",
      "                    changes_made=changes_made,\n",
      "                )\n",
      "            )\n",
      "            self.messages = self.messages[:-2]\n",
      "            new_function_pattern = r\"<new_function_names>\\s+(?P<new_function_names>.*?)</new_function_names>\"\n",
      "            new_function_matches = list(\n",
      "                re.finditer(new_function_pattern, extract_response, re.DOTALL)\n",
      "            )\n",
      "            for match_ in new_function_matches:\n",
      "                match = match_.groupdict()\n",
      "                new_function_matches = match[\"new_function_names\"]\n",
      "                new_function_matches = new_function_matches.split(\"\\n\")\n",
      "            extracted_pattern = r\"<<<<<<<\\s+EXTRACT\\s+(?P<updated_code>.*?)>>>>>>>\"\n",
      "            extract_matches = list(\n",
      "                re.finditer(extracted_pattern, extract_response, re.DOTALL)\n",
      "            )\n",
      "            new_code = None\n",
      "            for idx, match_ in enumerate(extract_matches[::-1]):\n",
      "                match = match_.groupdict()\n",
      "                updated_code = match[\"updated_code\"]\n",
      "                updated_code = updated_code.strip(\"\\n\")\n",
      "                if len(updated_code) < 150:  # too few characters\n",
      "                    continue\n",
      "                # too close to function length, just skip for now\n",
      "                if (\n",
      "                    len(updated_code) / (len(function_and_reference.function_code) + 1)\n",
      "                    > 0.9\n",
      "                ):\n",
      "                    continue\n",
      "                best_match = find_best_match(updated_code, recent_file_contents)\n",
      "                if best_match.score < 70:\n",
      "                    updated_code = \"\\n\".join(updated_code.split(\"\\n\")[1:])\n",
      "                    best_match = find_best_match(updated_code, recent_file_contents)\n",
      "                    if best_match.score < 80:\n",
      "                        updated_code = \"\\n\".join(updated_code.split(\"\\n\")[:-1])\n",
      "                        best_match = find_best_match(updated_code, recent_file_contents)\n",
      "                        if best_match.score < 80:\n",
      "                            continue\n",
      "                matched_lines = recent_file_contents.split(\"\\n\")[\n",
      "                    best_match.start : best_match.end\n",
      "                ]\n",
      "                # handle return edge case\n",
      "                if matched_lines[-1].strip().startswith(\"return\"):\n",
      "                    matched_lines = matched_lines[:-1]\n",
      "                extracted_original_code = \"\\n\".join(matched_lines)\n",
      "                extracted_exact_matches.append(extracted_original_code)\n",
      "        deduped_exact_matches = []\n",
      "        for extracted_exact_match in extracted_exact_matches:\n",
      "            if extracted_exact_match not in deduped_exact_matches:\n",
      "                deduped_exact_matches.append(extracted_exact_match)\n",
      "\n",
      "        new_function_names = []\n",
      "        existing_names = \", \".join(\n",
      "            [def_fn.name.strip(\"'\") for def_fn in all_defined_functions]\n",
      "        )\n",
      "        offset = 5\n",
      "        for idx in range(0, len(deduped_exact_matches), offset):\n",
      "            num_snippets = min(len(deduped_exact_matches), idx + offset) - idx\n",
      "            formatted_snippets = \"\\n\".join(\n",
      "                [\n",
      "                    f\"<function_to_name>\\n{snippet}\\n</function_to_name>\"\n",
      "                    for snippet in deduped_exact_matches[idx:idx+num_snippets]\n",
      "                ]\n",
      "            )\n",
      "            new_function_names.extend(NameBot(chat_logger=self.chat_logger).name_functions(\n",
      "                old_code=cloned_repo.get_file_contents(file_path),\n",
      "                snippets=formatted_snippets,\n",
      "                existing_names=existing_names,\n",
      "                count=num_snippets,\n",
      "            ))\n",
      "        for idx, extracted_original_code in enumerate(deduped_exact_matches):\n",
      "            if idx >= len(new_function_names):\n",
      "                break\n",
      "            new_code, _ = extract_method(\n",
      "                extracted_original_code,\n",
      "                file_path,\n",
      "                new_function_names[idx],\n",
      "                project_name=cloned_repo.repo_dir,\n",
      "            )\n",
      "        return new_code\n",
      "\n",
      "new_function_names = []\n",
      "        existing_names = \", \".join(\n",
      "            [def_fn.name.strip(\"'\") for def_fn in all_defined_functions]\n",
      "        )\n",
      "        offset = 5\n",
      "        for idx in range(0, len(deduped_exact_matches), offset):\n",
      "            num_snippets = min(len(deduped_exact_matches), idx + offset) - idx\n",
      "            formatted_snippets = \"\\n\".join(\n",
      "                [\n",
      "                    f\"<function_to_name>\\n{snippet}\\n</function_to_name>\"\n",
      "                    for snippet in deduped_exact_matches[idx:idx+num_snippets]\n",
      "                ]\n",
      "            )\n",
      "            new_function_names.extend(NameBot(chat_logger=self.chat_logger).name_functions(\n",
      "                old_code=cloned_repo.get_file_contents(file_path),\n",
      "                snippets=formatted_snippets,\n",
      "                existing_names=existing_names,\n",
      "                count=num_snippets,\n",
      "            ))\n",
      "        for idx, extracted_original_code in enumerate(deduped_exact_matches):\n",
      "            if idx >= len(new_function_names):\n",
      "                break\n",
      "9895 10896\n",
      "--- a/test.py\n",
      "+++ b/test.py\n",
      "@@ -268,3 +268,28 @@\n",
      "                 project_name=cloned_repo.repo_dir,\n",
      "             )\n",
      "         return new_code\n",
      "+\n",
      "+    def helper(self, all_defined_functions, deduped_exact_matches, cloned_repo, file_path, idx, extracted_original_code):\n",
      "+        new_function_names = []\n",
      "+        existing_names = \", \".join(\n",
      "+            [def_fn.name.strip(\"'\") for def_fn in all_defined_functions]\n",
      "+        )\n",
      "+        offset = 5\n",
      "+        for idx in range(0, len(deduped_exact_matches), offset):\n",
      "+            num_snippets = min(len(deduped_exact_matches), idx + offset) - idx\n",
      "+            formatted_snippets = \"\\n\".join(\n",
      "+                [\n",
      "+                    f\"<function_to_name>\\n{snippet}\\n</function_to_name>\"\n",
      "+                    for snippet in deduped_exact_matches[idx:idx+num_snippets]\n",
      "+                ]\n",
      "+            )\n",
      "+            new_function_names.extend(NameBot(chat_logger=self.chat_logger).name_functions(\n",
      "+                old_code=cloned_repo.get_file_contents(file_path),\n",
      "+                snippets=formatted_snippets,\n",
      "+                existing_names=existing_names,\n",
      "+                count=num_snippets,\n",
      "+            ))\n",
      "+        for idx, extracted_original_code in enumerate(deduped_exact_matches):\n",
      "+            if idx >= len(new_function_names):\n",
      "+                break\n",
      "+        return extracted_original_code, new_function_names, idx\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import rope.base.project\n",
    "from rope.refactor.extract import ExtractMethod\n",
    "from rope.refactor.inline import InlineMethod\n",
    "import re\n",
    "\n",
    "APOSTROPHE_MARKER = \"__APOSTROPHE__\"\n",
    "PERCENT_FORMAT_MARKER = \"__PERCENT_FORMAT__\"\n",
    "\n",
    "def serialize(text: str):\n",
    "    # Replace \"'{var}'\" with \"__APOSTROPHE__{var}__APOSTROPHE__\"\n",
    "    text = re.sub(r\"'{([^'}]*?)}'\", f\"{APOSTROPHE_MARKER}{{\\\\1}}{APOSTROPHE_MARKER}\", text)\n",
    "    # Replace \"%s\" with \"__PERCENT_FORMAT__\"\n",
    "    text = re.sub(r\"%\\((.*?)\\)s\", f\"{PERCENT_FORMAT_MARKER}{{\\\\1}}\", text)\n",
    "    return text\n",
    "\n",
    "def deserialize(text: str):\n",
    "    text = re.sub(f\"{APOSTROPHE_MARKER}{{(.*?)}}{APOSTROPHE_MARKER}\", \"'{\\\\1}'\", text)\n",
    "    text = re.sub(f\"{PERCENT_FORMAT_MARKER}{{(.*?)}}\", \"%(\\\\1)s\", text)\n",
    "    return text\n",
    "\n",
    "myproject = rope.base.project.Project('src')\n",
    "\n",
    "myresource = myproject.get_resource('test.py')\n",
    "contents = myresource.read()\n",
    "print(contents)\n",
    "serialized_contents = serialize(myresource.read())\n",
    "myresource.write(serialized_contents)\n",
    "extract_span = r\"\"\"new_function_names = []\n",
    "        existing_names = \", \".join(\n",
    "            [def_fn.name.strip(\"'\") for def_fn in all_defined_functions]\n",
    "        )\n",
    "        offset = 5\n",
    "        for idx in range(0, len(deduped_exact_matches), offset):\n",
    "            num_snippets = min(len(deduped_exact_matches), idx + offset) - idx\n",
    "            formatted_snippets = \"\\n\".join(\n",
    "                [\n",
    "                    f\"<function_to_name>\\n{snippet}\\n</function_to_name>\"\n",
    "                    for snippet in deduped_exact_matches[idx:idx+num_snippets]\n",
    "                ]\n",
    "            )\n",
    "            new_function_names.extend(NameBot(chat_logger=self.chat_logger).name_functions(\n",
    "                old_code=cloned_repo.get_file_contents(file_path),\n",
    "                snippets=formatted_snippets,\n",
    "                existing_names=existing_names,\n",
    "                count=num_snippets,\n",
    "            ))\n",
    "        for idx, extracted_original_code in enumerate(deduped_exact_matches):\n",
    "            if idx >= len(new_function_names):\n",
    "                break\"\"\"\n",
    "serialized_extract_span = serialize(extract_span)\n",
    "print(serialized_extract_span)\n",
    "\n",
    "start, end = serialized_contents.find(serialized_extract_span), serialized_contents.find(serialized_extract_span) + len(serialized_extract_span)\n",
    "print(start, end)\n",
    "\n",
    "try:\n",
    "    extractor = ExtractMethod(myproject, myresource, start, end)\n",
    "    change_set = extractor.get_changes(\"helper\", similar=True)\n",
    "    for change in change_set.changes:\n",
    "        if change.old_contents is not None:\n",
    "            change.old_contents = deserialize(change.old_contents)\n",
    "        else:\n",
    "            change.old_contents = deserialize(change.resource.read())\n",
    "        change.new_contents = deserialize(change.new_contents)\n",
    "    for change in change_set.changes:\n",
    "        print(change.get_description())\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "    raise e\n",
    "finally:\n",
    "    myresource.write(contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myresource = myproject.get_resource('mod3.py')\n",
    "contents = myresource.read()\n",
    "serialized_contents = serialize(myresource.read())\n",
    "myresource.write(serialized_contents)\n",
    "extract_span = r\"\"\"                openai.api_type = OPENAI_API_TYPE\n",
    "                openai.api_base = OPENAI_API_BASE\n",
    "                openai.api_version = OPENAI_API_VERSION\n",
    "                openai.api_key = AZURE_API_KEY\n",
    "                response = openai.ChatCompletion.create(\n",
    "                    engine=engine,\n",
    "                    model=model,\n",
    "                    messages=messages,\n",
    "                    max_tokens=max_tokens,\n",
    "                    temperature=temperature,\n",
    "                    timeout=OPENAI_TIMEOUT,\n",
    "                )\"\"\"\n",
    "serialized_extract_span = serialize(extract_span)\n",
    "print(serialized_extract_span)\n",
    "\n",
    "start, end = serialized_contents.find(serialized_extract_span), serialized_contents.find(serialized_extract_span) + len(serialized_extract_span)\n",
    "print(start, end)\n",
    "\n",
    "try:\n",
    "    extractor = ExtractMethod(myproject, myresource, start, end)\n",
    "    change_set = extractor.get_changes(\"helper\", similar=True)\n",
    "    for change in change_set.changes:\n",
    "        if change.old_contents is not None:\n",
    "            change.old_contents = deserialize(change.old_contents)\n",
    "        else:\n",
    "            change.old_contents = deserialize(change.resource.read())\n",
    "        change.new_contents = deserialize(change.new_contents)\n",
    "    for change in change_set.changes:\n",
    "        print(change.get_description())\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "finally:\n",
    "    myresource.write(contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "85"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(change.new_contents) - len(change.old_contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restructuring <\n",
      "openai.api_key = ${api_key}\n",
      "openai.api_base = ${api_base}\n",
      "openai.api_version = ${api_version}\n",
      "openai.api_type = ${api_type}\n",
      "logger.info(f\"Calling {model} with OpenAI.\")\n",
      "response = openai.ChatCompletion.create(\n",
      "    model=${model},\n",
      "    messages=${messages},\n",
      "    max_tokens=${max_tokens},\n",
      "    temperature=${temperature},\n",
      "    timeout=${timeout},\n",
      "    seed=${seed},\n",
      ")> to <print(\n",
      "    ${api_key},\n",
      "    ${api_base},\n",
      "    ${api_version},\n",
      "    ${api_type},\n",
      "    ${model},\n",
      "    ${messages},\n",
      "    ${max_tokens},\n",
      "    ${temperature},\n",
      "    ${timeout},\n",
      "    ${seed},\n",
      ")>:\n",
      "\n",
      "\n",
      "--- a/mod3.py\n",
      "+++ b/mod3.py\n",
      "@@ -29,18 +29,17 @@\n",
      "             ):\n",
      "                 engine = OPENAI_API_ENGINE_GPT4_32K\n",
      "             if OPENAI_API_TYPE is None or engine is None:\n",
      "-                openai.api_key = OPENAI_API_KEY\n",
      "-                openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                openai.api_version = None\n",
      "-                openai.api_type = \"open_ai\"\n",
      "-                logger.info(f\"Calling {model} with OpenAI.\")\n",
      "-                response = openai.ChatCompletion.create(\n",
      "-                    model=model,\n",
      "-                    messages=messages,\n",
      "-                    max_tokens=max_tokens,\n",
      "-                    temperature=temperature,\n",
      "-                    timeout=OPENAI_TIMEOUT,\n",
      "-                    seed=SEED,\n",
      "+                print(\n",
      "+                    OPENAI_API_KEY,\n",
      "+                    \"https://api.openai.com/v1\",\n",
      "+                    None,\n",
      "+                    \"open_ai\",\n",
      "+                    model,\n",
      "+                    messages,\n",
      "+                    max_tokens,\n",
      "+                    temperature,\n",
      "+                    OPENAI_TIMEOUT,\n",
      "+                    SEED,\n",
      "                 )\n",
      "                 return response[\"choices\"][0].message.content\n",
      "             # validity checks for MULTI_REGION_CONFIG\n",
      "@@ -100,18 +99,17 @@\n",
      "         except Exception as e:\n",
      "             if OPENAI_API_KEY:\n",
      "                 try:\n",
      "-                    openai.api_key = OPENAI_API_KEY\n",
      "-                    openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                    openai.api_version = None\n",
      "-                    openai.api_type = \"open_ai\"\n",
      "-                    logger.info(f\"Calling {model} with OpenAI.\")\n",
      "-                    response = openai.ChatCompletion.create(\n",
      "-                        model=model,\n",
      "-                        messages=messages,\n",
      "-                        max_tokens=max_tokens,\n",
      "-                        temperature=temperature,\n",
      "-                        timeout=OPENAI_TIMEOUT,\n",
      "-                        seed=SEED,\n",
      "+                    print(\n",
      "+                        OPENAI_API_KEY,\n",
      "+                        \"https://api.openai.com/v1\",\n",
      "+                        None,\n",
      "+                        \"open_ai\",\n",
      "+                        model,\n",
      "+                        messages,\n",
      "+                        max_tokens,\n",
      "+                        temperature,\n",
      "+                        OPENAI_TIMEOUT,\n",
      "+                        SEED,\n",
      "                     )\n",
      "                     return response[\"choices\"][0].message.content\n",
      "                 except SystemExit:\n",
      "\n",
      "--- a/test.py\n",
      "+++ b/test.py\n",
      "@@ -60,18 +60,17 @@\n",
      "             ):\n",
      "                 engine = OPENAI_API_ENGINE_GPT4_32K\n",
      "             if OPENAI_API_TYPE is None or engine is None:\n",
      "-                openai.api_key = OPENAI_API_KEY\n",
      "-                openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                openai.api_version = None\n",
      "-                openai.api_type = \"open_ai\"\n",
      "-                logger.info(f\"Calling {model} with OpenAI.\")\n",
      "-                response = openai.ChatCompletion.create(\n",
      "-                    model=model,\n",
      "-                    messages=messages,\n",
      "-                    max_tokens=max_tokens,\n",
      "-                    temperature=temperature,\n",
      "-                    timeout=OPENAI_TIMEOUT,\n",
      "-                    seed=SEED,\n",
      "+                print(\n",
      "+                    OPENAI_API_KEY,\n",
      "+                    \"https://api.openai.com/v1\",\n",
      "+                    None,\n",
      "+                    \"open_ai\",\n",
      "+                    model,\n",
      "+                    messages,\n",
      "+                    max_tokens,\n",
      "+                    temperature,\n",
      "+                    OPENAI_TIMEOUT,\n",
      "+                    SEED,\n",
      "                 )\n",
      "                 return response[\"choices\"][0].message.content\n",
      "             # validity checks for MULTI_REGION_CONFIG\n",
      "@@ -131,18 +130,17 @@\n",
      "         except Exception as e:\n",
      "             if OPENAI_API_KEY:\n",
      "                 try:\n",
      "-                    openai.api_key = OPENAI_API_KEY\n",
      "-                    openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                    openai.api_version = None\n",
      "-                    openai.api_type = \"open_ai\"\n",
      "-                    logger.info(f\"Calling {model} with OpenAI.\")\n",
      "-                    response = openai.ChatCompletion.create(\n",
      "-                        model=model,\n",
      "-                        messages=messages,\n",
      "-                        max_tokens=max_tokens,\n",
      "-                        temperature=temperature,\n",
      "-                        timeout=OPENAI_TIMEOUT,\n",
      "-                        seed=SEED,\n",
      "+                    print(\n",
      "+                        OPENAI_API_KEY,\n",
      "+                        \"https://api.openai.com/v1\",\n",
      "+                        None,\n",
      "+                        \"open_ai\",\n",
      "+                        model,\n",
      "+                        messages,\n",
      "+                        max_tokens,\n",
      "+                        temperature,\n",
      "+                        OPENAI_TIMEOUT,\n",
      "+                        SEED,\n",
      "                     )\n",
      "                     return response[\"choices\"][0].message.content\n",
      "                 except SystemExit:\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "myproject = rope.base.project.Project('./src')\n",
    "mod1 = myproject.get_resource('mod3.py')\n",
    "mod1.write(r\"\"\"\n",
    "class OpenAIProxy:\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    @file_cache(ignore_params=[])\n",
    "    def call_openai(self, model, messages, max_tokens, temperature) -> str:\n",
    "        try:\n",
    "            engine = None\n",
    "            if model in OPENAI_EXCLUSIVE_MODELS and OPENAI_API_TYPE != \"azure\":\n",
    "                logger.info(f\"Calling OpenAI exclusive model. {model}\")\n",
    "                raise Exception(\"OpenAI exclusive model.\")\n",
    "            if (\n",
    "                model == \"gpt-3.5-turbo-16k\"\n",
    "                or model == \"gpt-3.5-turbo-16k-0613\"\n",
    "                and OPENAI_API_ENGINE_GPT35 is not None\n",
    "            ):\n",
    "                engine = OPENAI_API_ENGINE_GPT35\n",
    "            elif (\n",
    "                model == \"gpt-4\"\n",
    "                or model == \"gpt-4-0613\"\n",
    "                and OPENAI_API_ENGINE_GPT4 is not None\n",
    "            ):\n",
    "                engine = OPENAI_API_ENGINE_GPT4\n",
    "            elif (\n",
    "                model == \"gpt-4-32k\"\n",
    "                or model == \"gpt-4-32k-0613\"\n",
    "                and OPENAI_API_ENGINE_GPT4_32K is not None\n",
    "            ):\n",
    "                engine = OPENAI_API_ENGINE_GPT4_32K\n",
    "            if OPENAI_API_TYPE is None or engine is None:\n",
    "                openai.api_key = OPENAI_API_KEY\n",
    "                openai.api_base = \"https://api.openai.com/v1\"\n",
    "                openai.api_version = None\n",
    "                openai.api_type = \"open_ai\"\n",
    "                response = openai.ChatCompletion.create(\n",
    "                    model=model,\n",
    "                    messages=messages,\n",
    "                    max_tokens=max_tokens,\n",
    "                    temperature=temperature,\n",
    "                    timeout=OPENAI_TIMEOUT,\n",
    "                    seed=SEED,\n",
    "                )\n",
    "                return response[\"choices\"][0].message.content\n",
    "            # validity checks for MULTI_REGION_CONFIG\n",
    "            if (\n",
    "                MULTI_REGION_CONFIG is None\n",
    "                or not isinstance(MULTI_REGION_CONFIG, list)\n",
    "                or len(MULTI_REGION_CONFIG) == 0\n",
    "                or not isinstance(MULTI_REGION_CONFIG[0], list)\n",
    "            ):\n",
    "                logger.info(\n",
    "                    f\"Calling {model} with engine {engine} on Azure url {OPENAI_API_BASE}.\"\n",
    "                )\n",
    "                openai.api_type = OPENAI_API_TYPE\n",
    "                openai.api_base = OPENAI_API_BASE\n",
    "                openai.api_version = OPENAI_API_VERSION\n",
    "                openai.api_key = AZURE_API_KEY\n",
    "                response = openai.ChatCompletion.create(\n",
    "                    engine=engine,\n",
    "                    model=model,\n",
    "                    messages=messages,\n",
    "                    max_tokens=max_tokens,\n",
    "                    temperature=temperature,\n",
    "                    timeout=OPENAI_TIMEOUT,\n",
    "                )\n",
    "                return response[\"choices\"][0].message.content\n",
    "            # multi region config is a list of tuples of (region_url, api_key)\n",
    "            # we will try each region in order until we get a response\n",
    "            # randomize the order of the list\n",
    "            SHUFFLED_MULTI_REGION_CONFIG = random.sample(\n",
    "                MULTI_REGION_CONFIG, len(MULTI_REGION_CONFIG)\n",
    "            )\n",
    "            for region_url, api_key in SHUFFLED_MULTI_REGION_CONFIG:\n",
    "                try:\n",
    "                    logger.info(\n",
    "                        f\"Calling {model} with engine {engine} on Azure url {region_url}.\"\n",
    "                    )\n",
    "                    openai.api_key = api_key\n",
    "                    openai.api_base = region_url\n",
    "                    openai.api_version = OPENAI_API_VERSION\n",
    "                    openai.api_type = OPENAI_API_TYPE\n",
    "                    response = openai.ChatCompletion.create(\n",
    "                        engine=engine,\n",
    "                        model=model,\n",
    "                        messages=messages,\n",
    "                        max_tokens=max_tokens,\n",
    "                        temperature=temperature,\n",
    "                        timeout=OPENAI_TIMEOUT,\n",
    "                    )\n",
    "                    return response[\"choices\"][0].message.content\n",
    "                except SystemExit:\n",
    "                    raise SystemExit\n",
    "                except Exception as e:\n",
    "                    logger.exception(f\"Error calling {region_url}: {e}\")\n",
    "            raise Exception(\"No Azure regions available\")\n",
    "        except SystemExit:\n",
    "            raise SystemExit\n",
    "        except Exception as e:\n",
    "            if OPENAI_API_KEY:\n",
    "                try:\n",
    "                    openai.api_key = OPENAI_API_KEY\n",
    "                    openai.api_base = \"https://api.openai.com/v1\"\n",
    "                    openai.api_version = None\n",
    "                    openai.api_type = \"open_ai\"\n",
    "                    response = openai.ChatCompletion.create(\n",
    "                        model=model,\n",
    "                        messages=messages,\n",
    "                        max_tokens=max_tokens,\n",
    "                        temperature=temperature,\n",
    "                        timeout=OPENAI_TIMEOUT,\n",
    "                        seed=SEED,\n",
    "                    )\n",
    "                    return response[\"choices\"][0].message.content\n",
    "                except SystemExit:\n",
    "                    raise SystemExit\n",
    "                except Exception as _e:\n",
    "                    logger.error(f\"OpenAI API Key found but error: {_e}\")\n",
    "            logger.error(f\"OpenAI API Key not found and Azure Error: {e}\")\n",
    "            # Raise exception to report error\n",
    "            raise e\n",
    "\"\"\")\n",
    "\n",
    "from rope.refactor import restructure\n",
    "pattern = r\"\"\"\n",
    "openai.api_key = ${api_key}\n",
    "openai.api_base = ${api_base}\n",
    "openai.api_version = ${api_version}\n",
    "openai.api_type = ${api_type}\n",
    "response = openai.ChatCompletion.create(\n",
    "    model=${model},\n",
    "    messages=${messages},\n",
    "    max_tokens=${max_tokens},\n",
    "    temperature=${temperature},\n",
    "    timeout=${timeout},\n",
    "    seed=${seed},\n",
    ")\"\"\"\n",
    "goal = \"\"\"print(\n",
    "    ${api_key},\n",
    "    ${api_base},\n",
    "    ${api_version},\n",
    "    ${api_type},\n",
    "    ${model},\n",
    "    ${messages},\n",
    "    ${max_tokens},\n",
    "    ${temperature},\n",
    "    ${timeout},\n",
    "    ${seed},\n",
    ")\"\"\"\n",
    "\n",
    "restructuring = restructure.Restructure(myproject, pattern, goal)\n",
    "print(restructuring.get_changes().get_description())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using function <_call_openai>:\n",
      "\n",
      "\n",
      "--- a/utils/openai_proxy.py\n",
      "+++ b/utils/openai_proxy.py\n",
      "@@ -60,18 +60,7 @@\n",
      "             ):\n",
      "                 engine = OPENAI_API_ENGINE_GPT4_32K\n",
      "             if OPENAI_API_TYPE is None or engine is None:\n",
      "-                openai.api_key = OPENAI_API_KEY\n",
      "-                openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                openai.api_version = None\n",
      "-                openai.api_type = \"open_ai\"\n",
      "-                response = openai.ChatCompletion.create(\n",
      "-                    model=model,\n",
      "-                    messages=messages,\n",
      "-                    max_tokens=max_tokens,\n",
      "-                    temperature=temperature,\n",
      "-                    timeout=OPENAI_TIMEOUT,\n",
      "-                    seed=SEED,\n",
      "-                )\n",
      "+                _call_openai(model, messages, max_tokens, temperature, OPENAI_API_KEY, \"https://api.openai.com/v1\", None, \"open_ai\")\n",
      "                 return response[\"choices\"][0].message.content\n",
      "             # validity checks for MULTI_REGION_CONFIG\n",
      "             if (\n",
      "@@ -130,18 +119,7 @@\n",
      "         except Exception as e:\n",
      "             if OPENAI_API_KEY:\n",
      "                 try:\n",
      "-                    openai.api_key = OPENAI_API_KEY\n",
      "-                    openai.api_base = \"https://api.openai.com/v1\"\n",
      "-                    openai.api_version = None\n",
      "-                    openai.api_type = \"open_ai\"\n",
      "-                    response = openai.ChatCompletion.create(\n",
      "-                        model=model,\n",
      "-                        messages=messages,\n",
      "-                        max_tokens=max_tokens,\n",
      "-                        temperature=temperature,\n",
      "-                        timeout=OPENAI_TIMEOUT,\n",
      "-                        seed=SEED,\n",
      "-                    )\n",
      "+                    _call_openai(model, messages, max_tokens, temperature, OPENAI_API_KEY, \"https://api.openai.com/v1\", None, \"open_ai\")\n",
      "                     return response[\"choices\"][0].message.content\n",
      "                 except SystemExit:\n",
      "                     raise SystemExit\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from rope.refactor import usefunction\n",
    "\n",
    "sweep_project = rope.base.project.Project('../../sweepai', ignored_resources=[\"sandbox\"])\n",
    "resource = sweep_project.get_resource('utils/str_utils.py')\n",
    "# resource = sweep_project.get_resource('utils/openai_proxy.py')\n",
    "# start = resource.read().find('_call_openai')\n",
    "start = resource.read().find('entities_split')\n",
    "use_function = usefunction.UseFunction(sweep_project, resource, start)\n",
    "print(use_function.get_changes().get_description())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "from rope.base.project import Project\n",
    "from rope.contrib.autoimport import AutoImport\n",
    "\n",
    "project = Project(\"../../sweepai\")\n",
    "autoimport = AutoImport(project)\n",
    "# autoimport.generate_resource_cache()  # Generates a cache of the local modules, from the project you're working on\n",
    "autoimport.generate_modules_cache([\"../../sweepai\", \".\"])  # Generates a cache of external modules\n",
    "print(autoimport.import_assist(\"sweepai\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 ('sweepai-u_CIt3kb-py3.10')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "25d341f3248a096a89b9dbf6eec8e41f63aed02f6ba059df22a49224e3e8f1b0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
